import base64

import gradio as gr
with open("./实验室.png", "rb") as image_file:
    encoded_string = base64.b64encode(image_file.read()).decode("utf-8")


# 假设以下函数已定义
def do_decode(input_wav_path, input_prompt):
    # 模拟处理逻辑，返回处理结果
    return "处理结果示例"

def download_audio(input_wav_path):
    # 模拟下载逻辑
    return input_wav_path

def save_to_jsonl(if_correct, wav, prompt, res):
    # 模拟保存逻辑
    pass

TASK_PROMPT_MAPPING = {
    "ASR (Automatic Speech Recognition)": "ASR任务提示",
    # 其他任务映射
}

def get_wav_from_token_list(input_list):
    return "./input_data/question1.wav"

# 创建Gradio界面
with gr.Blocks() as demo:
    # 添加标题
    gr.Markdown(
        f"""
        <div style="display: flex; align-items: center; justify-content: center; text-align: center;">
            <h1 style="font-family: 'Arial', sans-serif; color: #014377; font-size: 32px; margin-bottom: 0; display: inline-block; vertical-align: middle;">
                OSUM Speech Understanding Model Test
            </h1>
        </div>
        """
    )

    # 添加音频输入和任务选择
    with gr.Row():
        with gr.Column(scale=1, min_width=300):
            # audio_input = gr.Audio(label="录音", type="filepath")
            audio_input = gr.Audio(label="录音", sources=["microphone", "upload"], type="filepath", visible=True)
        with gr.Column(scale=1, min_width=300):  # 给输出框设置最小宽度，确保等高对齐
            output_text = gr.Textbox(label="输出结果", lines=6, placeholder="生成的结果将显示在这里...",
                                     interactive=False)

    # 添加任务选择和自定义输入框
    with gr.Row():
        task_dropdown = gr.Dropdown(
            label="任务",
            choices=list(TASK_PROMPT_MAPPING.keys()) + ["自主输入文本", "TTS任务"],  # 新增"TTS任务"选项
            value="ASR (Automatic Speech Recognition)"
        )
        custom_prompt_input = gr.Textbox(label="自定义任务提示", placeholder="请输入自定义任务提示...",
                                         visible=False)  # 新增文本输入框
        tts_input = gr.Textbox(label="TTS输入文本", placeholder="请输入TTS任务的文本...", visible=False)  # 新增TTS输入框

    # 添加音频播放组件
    audio_player = gr.Audio(label="播放音频", type="filepath", interactive=False)

    # 添加按钮（下载按钮在左边，开始处理按钮在右边）
    with gr.Row():
        download_button = gr.DownloadButton("下载音频", variant="secondary",
                                            elem_classes=["button-height", "download-button"])
        submit_button = gr.Button("开始处理", variant="primary", elem_classes=["button-height", "submit-button"])

    # 添加确认组件
    with gr.Row(visible=False) as confirmation_row:
        gr.Markdown("请判断结果是否正确：")
        confirmation_buttons = gr.Radio(
            choices=["正确", "错误"],
            label="",
            interactive=True,
            container=False,
            elem_classes="confirmation-buttons"
        )
        save_button = gr.Button("提交反馈", variant="secondary")

    # 添加底部内容
    with gr.Row():
        # 底部内容容器
        with gr.Column(scale=1, min_width=800):  # 设置最小宽度以确保内容居中
            gr.Markdown(
                f"""
                <div style="position: fixed; bottom: 20px; left: 50%; transform: translateX(-50%); display: flex; align-items: center; justify-content: center; gap: 20px;">
                    <div style="text-align: center;">
                        <p style="margin: 0;"><strong>Audio, Speech and Language Processing Group (ASLP@NPU),</strong></p>
                        <p style="margin: 0;"><strong>Northwestern Polytechnical University</strong></p>
                    </div>
                    <img src="data:image/png;base64,{encoded_string}" alt="OSUM Logo" style="height: 80px; width: auto;">
                </div>
                """
            )

    # 绑定事件
    def show_confirmation(output_res, input_wav_path, input_prompt):
        return gr.update(visible=True), output_res, input_wav_path, input_prompt

    def save_result(if_correct, wav, prompt, res):
        save_to_jsonl(if_correct, wav, prompt, res)
        return gr.update(visible=False)

    def handle_submit(input_wav_path, task_choice, custom_prompt, tts_text):
        if task_choice == "自主输入文本":
            input_prompt = custom_prompt + "_self_prompt"  # 自定义输入加上"_self_prompt"
        elif task_choice == "TTS任务":
            input_prompt = tts_text + "_TTS"  # TTS输入加上"_TTS"
        else:
            input_prompt = TASK_PROMPT_MAPPING.get(task_choice, "未知任务类型")  # 使用预定义的提示
        output_res = do_decode(input_wav_path, input_prompt)
        wav_path_output = input_wav_path
        if task_choice == "TTS任务":
            # assert isinstance(output_res, list), "TTS任务的输出必须为列表"
            wav_path_output = get_wav_from_token_list(output_res)
        return output_res, wav_path_output

    task_dropdown.change(
        fn=lambda choice: gr.update(visible=choice == "自主输入文本"),
        inputs=task_dropdown,
        outputs=custom_prompt_input
    )

    task_dropdown.change(
        fn=lambda choice: gr.update(visible=choice == "TTS任务"),
        inputs=task_dropdown,
        outputs=tts_input
    )

    submit_button.click(
        fn=handle_submit,
        inputs=[audio_input, task_dropdown, custom_prompt_input, tts_input],
        outputs=[output_text, audio_player]
    ).then(
        fn=show_confirmation,
        inputs=[output_text, audio_input, task_dropdown],
        outputs=[confirmation_row, output_text, audio_input, task_dropdown]
    )

    download_button.click(
        fn=download_audio,
        inputs=[audio_input],
        outputs=[download_button]  # 输出到 download_button
    )

    save_button.click(
        fn=save_result,
        inputs=[confirmation_buttons, audio_input, task_dropdown, output_text],
        outputs=confirmation_row
    )

# 启动Gradio界面
demo.launch()